#include "CudaRender.h"
#include "image.h"
#include "sceneLoader.h"
#include "noise.h"
#include "../include/helper_cuda.h"

#include <cuda.h>
#include <cuda_runtime.h>
#include <cuda_fp16.h>

#include <iostream>
#include <string>

// Inline Functino
#include "lookupColor.cu_inl"
#include "noiseCuda.cu_inl"
#include "CudaUtil.cu_inl"

// Cuda Global Struct
struct GlobalConstStruct
{
    SceneName sceneName;

    int num_Circles;

    // Cuda Global Memory Data pointer to global memory
    float *position; // 3 * num_circles * sizeof(float)
    float *velocity; // 3 * num_circles * sizeof(float)
    float *color;    // 3 * num_circles * sizeof(float)
    float *radius;   // 3 * num_circles * sizeof(float)

    int imageWidth;
    int imageHeight;
    float *imageData; // 4 * num_circles * sizeof(float)
};


// cuda Const Render Params
__constant__ GlobalConstStruct cuConstRenderParams;

// cuda Const Noise
__constant__ int cuConstNoiseYPermutationTable[256];
__constant__ int cuConstNoiseXPermutationTable[256];
__constant__ float cuConstNoise1DValueTable[256];

#define COLOR_MAP_SIZE 5
__constant__ float cuConstColorRamp[COLOR_MAP_SIZE][3];

// Kernel Functions

__global__ void kernelClearImageSnowflake()
{

    int imageX = blockIdx.x * blockDim.x + threadIdx.x;
    int imageY = blockIdx.y * blockDim.y + threadIdx.y;

    int width = cuConstRenderParams.imageWidth;
    int height = cuConstRenderParams.imageHeight;

    if (imageX >= width || imageY >= height)
        return;

    int offset = 4 * (imageY * width + imageX);
    float shade = .4f + .45f * static_cast<float>(height - imageY) / height;
    float4 value = make_float4(shade, shade, shade, 1.f);

    // write to global memory: As an optimization use a float4
    *(float4 *)(&cuConstRenderParams.imageData[offset]) = value;
}

__global__ void kernelClearImage(float r, float g, float b, float a)
{

    int imageX = blockIdx.x * blockDim.x + threadIdx.x;
    int imageY = blockIdx.y * blockDim.y + threadIdx.y;

    int width = cuConstRenderParams.imageWidth;
    int height = cuConstRenderParams.imageHeight;

    if (imageX >= width || imageY >= height)
        return;

    int offset = 4 * (imageY * width + imageX);
    float4 value = make_float4(r, g, b, a);

    // write to global memory: As an optimization use a float4
    *(float4 *)(&cuConstRenderParams.imageData[offset]) = value;
}

__global__ void kernelAdvanceSnowflake()
{
    int index = blockDim.x * blockIdx.x + threadIdx.x;

    if (index >= cuConstRenderParams.num_Circles)
        return;

    const float dt = 1.0f / 100.0f;
    const float kGravity = -1.8f;
    const float kDragCoeff = 2.0f;
    const float snow_mass = 0.1f;

    int index3 = index * 3;

    float *positionPtr = &cuConstRenderParams.position[index3];
    float *velocityPtr = &cuConstRenderParams.velocity[index3];

    float3 position = *(float3 *)(positionPtr);
    float3 velocity = *(float3 *)(velocityPtr);

    float forceScaling = fmin(fmax(1.0f - position.z, 0.1f), 1.0f);

    float3 noiseInput;
    noiseInput.x = 10.0f * position.x;
    noiseInput.y = 10.0f * position.y;
    noiseInput.z = 255.0f * position.z;
    float2 noiseForce = cudaVec2CellNoise(noiseInput, index);

    float2 dragForce;
    dragForce.x = -100.0f * kDragCoeff * velocity.x;
    dragForce.y = -1.0f * kDragCoeff * velocity.y;

    // update positions
    position.x += velocity.x * dt;
    position.y += velocity.y * dt;

    // update velocities v = v0 + at -> v = v0 + a * dt
    // F = ma -> a = F / m -> a = F * m_inv -> Assum mass_snow = 0.1kg -> m_inv = 10
    velocity.x += forceScaling * snow_mass * (noiseForce.x + dragForce.y) * dt;
    velocity.y += forceScaling * snow_mass * (noiseForce.y + dragForce.y) * dt;

    float radius = cuConstRenderParams.radius[index];

    if ((position.y + radius) < 0.0f || (position.x + radius) < -0.0f || (position.x - radius) > 1.0f)
    {
        noiseInput.x = 255.0f * position.x;
        noiseInput.y = 255.0f * position.y;
        noiseInput.z = 255.0f * position.z;
        noiseForce = cudaVec2CellNoise(noiseInput, index);

        position.x = 0.5f + 0.5f * noiseForce.x;
        position.y = 1.35f + radius;

        velocity.x = 2.0f * noiseForce.y;
        velocity.y = 0.0f;
    }

    *(float3 *)(positionPtr) = position;
    *(float3 *)(velocityPtr) = velocity;
}

__global__ void kernelAdvanceBouncingBalls()
{
    const float dt = 1.0f / 60.0f;
    const float kGravity = -1.8f;
    const float kDragCoeff = -0.8f;
    const float epsilon = 0.0001;

    int index = blockDim.x * blockIdx.x + threadIdx.x;

    if (index > cuConstRenderParams.num_Circles)
        return;

    float *velocity = cuConstRenderParams.velocity;
    float *position = cuConstRenderParams.position;

    int index3 = index * 3;
    float oldVelocity_y = cuConstRenderParams.velocity[index3 + 1];
    float oldPosition_y = cuConstRenderParams.position[index3 + 1];

    if (oldVelocity_y == 0.0f && oldPosition_y == 0.0f)
        return;

    if (oldPosition_y < 0.0f && oldVelocity_y < 0.0f)
    {
        // Bouncing Ball Note : kDrageCoeff is negative
        velocity[index3 + 1] *= kDragCoeff;
    }

    // else -> Update Ball velocity and position
    velocity[index3 + 1] += kGravity * dt;

    position[index3 + 1] += velocity[index3 + 1] * dt;

    if (fabsf(velocity[index3 + 1] - oldVelocity_y) < epsilon && oldPosition_y < 0.0f && fabsf(position[index3 + 1] - oldPosition_y) < epsilon)
    {
        // stop ball
        velocity[index3 + 1] = 0.f;
        position[index3 + 1] = 0.f;

        // if wants re-simulation
        // velocity[index3 + 1] = 0.1f;
        // position[index3 + 1] = 1.0f; // y head is 1.0f
    }
}

__global__ void kernelAdvanceHypnosis()
{
    int index = blockIdx.x * blockDim.x + threadIdx.x;
    if (index >= cuConstRenderParams.num_Circles)
        return;

    float *radius = cuConstRenderParams.radius;

    float cutOff = 0.5f;
    // place circle back in center after reaching threshold radisus
    if (radius[index] > cutOff)
    {
        radius[index] = 0.02f;
    }
    else
    {
        radius[index] += 0.01f;
    }
}

#define NUM_FIREWORKS 15
#define NUM_SPARKS 20

__global__ void kernelAdvanceFireWorks()
{
    const float dt = 1.f / 60.f;
    const float pi = 3.14159;
    const float maxDist = 0.25f;

    float *velocity = cuConstRenderParams.velocity;
    float *position = cuConstRenderParams.position;
    float *radius = cuConstRenderParams.radius;

    int index = blockIdx.x * blockDim.x + threadIdx.x;
    if (index >= cuConstRenderParams.num_Circles)
        return;

    if (0 <= index && index < NUM_FIREWORKS)
    {
        // firework center; no update
        return;
    }

    // determine the fire-work center/spark indices
    int fIdx = (index - NUM_FIREWORKS) / NUM_SPARKS;
    int sfIdx = (index - NUM_FIREWORKS) % NUM_SPARKS;

    int index3i = 3 * fIdx;
    int sIdx = NUM_FIREWORKS + fIdx * NUM_SPARKS + sfIdx;
    int index3j = 3 * sIdx;

    float cx = position[index3i];
    float cy = position[index3i + 1];

    // update position
    position[index3j] += velocity[index3j] * dt;
    position[index3j + 1] += velocity[index3j + 1] * dt;

    // fire-work sparks
    float sx = position[index3j];
    float sy = position[index3j + 1];

    // compute vector from firework-spark
    float cxsx = sx - cx;
    float cysy = sy - cy;

    // compute distance from fire-work
    float dist = sqrt(cxsx * cxsx + cysy * cysy);
    if (dist > maxDist)
    { // restore to starting position
        // random starting position on fire-work's rim
        float angle = (sfIdx * 2 * pi) / NUM_SPARKS;
        float sinA = sin(angle);
        float cosA = cos(angle);
        float x = cosA * radius[fIdx];
        float y = sinA * radius[fIdx];

        position[index3j] = position[index3i] + x;
        position[index3j + 1] = position[index3i + 1] + y;
        position[index3j + 2] = 0.0f;

        // travel scaled unit length
        velocity[index3j] = cosA / 5.0;
        velocity[index3j + 1] = sinA / 5.0;
        velocity[index3j + 2] = 0.0f;
    }
}

// --------------------------------- Wait to Implementation ---------------------------
__device__ void shadePixel(int circleIndex, float2 pixelCenter, float3 p, float4 *imagePtr)
{
    float diffX = p.x - pixelCenter.x;
    float diffY = p.y - pixelCenter.y;
    float pixelDist = diffX * diffX + diffY * diffY;

    float rad = cuConstRenderParams.radius[circleIndex];
    float maxDist = rad * rad;

    // circle does not contribute to the image
    if (pixelDist > maxDist)
        return;

    float3 rgb;
    float alpha;
}

__global__ void kernelRenderCircles()
{
    int index = blockDim.x * blockIdx.x + threadIdx.x;

    if (index > cuConstRenderParams.num_Circles)
        return;

    int index3 = index * 3;

    float3 p = *(float3 *)(&cuConstRenderParams.position[index3]);
    float rad = cuConstRenderParams.radius[index];

    short imageWidth = cuConstRenderParams.imageWidth;
    short imageHeight = cuConstRenderParams.imageHeight;
    short minX = static_cast<short>(imageWidth * (p.x - rad));
    short maxX = static_cast<short>(imageWidth * (p.x + rad)) + 1;
    short minY = static_cast<short>(imageHeight * (p.y - rad));
    short maxY = static_cast<short>(imageHeight * (p.y + rad)) + 1;

    short screenMinX = Clamp(minX, short(0.0f), imageWidth);
    short screenMaxX = Clamp(maxX, short(0.0f), imageWidth);
    short screenMinY = Clamp(minY, short(0.0f), imageWidth);
    short screenMaxY = Clamp(maxY, short(0.0f), imageWidth);

    float invWidth = 1.f / imageWidth;
    float invHeight = 1.f / imageHeight;

    // for all pixels in the bonding box
    for (int pixelY = screenMinY; pixelY < screenMaxY; pixelY++)
    {
        float4 *imgPtr = (float4 *)(&cuConstRenderParams.imageData[4 * (pixelY * imageWidth + screenMinX)]);
        for (int pixelX = screenMinX; pixelX < screenMaxX; pixelX++)
        {
            float2 pixelCenterNorm = make_float2(invWidth * (static_cast<float>(pixelX) + 0.5f),
                                                 invHeight * (static_cast<float>(pixelY) + 0.5f));
            shadePixel(index, pixelCenterNorm, p, imgPtr);
            imgPtr++;
        }
    }
}

__device__ void shadePixel_Hamster(float4 *imagePtr, int circle_index, int pixelDist, float3 p)
{
    float3 rgb;
    float alpha;

    float rad = cuConstRendererParams.radius[circle_index];

    if (cuConstRendererParams.sceneName == SNOWFLAKES || cuConstRendererParams.sceneName == SNOWFLAKES_SINGLE_FRAME)
    {

        const float kCircleMaxAlpha = .5f;
        const float falloffScale = 4.f;

        float normPixelDist = rsqrtf(pixelDist) / rad;
        rgb = lookupColor(normPixelDist);

        float maxAlpha = .6f + .4f * (1.f - p.z);
        maxAlpha = kCircleMaxAlpha * fmaxf(fminf(maxAlpha, 1.f), 0.f); // kCircleMaxAlpha * clamped value
        alpha = maxAlpha * exp(-1.f * falloffScale * normPixelDist * normPixelDist);
    }
    else
    {
        // simple: each circle has an assigned color
        int index3 = 3 * circle_index;
        rgb = *(float3 *)&(cuConstRendererParams.color[index3]);
        alpha = .5f;
    }

    float oneMinusAlpha = 1.f - alpha;

    float4 existingColor = *imagePtr;
    float4 newColor;
    newColor.x = alpha * rgb.x + oneMinusAlpha * existingColor.x;
    newColor.y = alpha * rgb.y + oneMinusAlpha * existingColor.y;
    newColor.z = alpha * rgb.z + oneMinusAlpha * existingColor.z;
    newColor.w = alpha + existingColor.w;

    *imagePtr = newColor;
}

// Goal is Render Every Pixel in Image
#include "circleBoxTest.cu_inl"

__global__ void kernelRenderPixel()
{
    int pixel_location_i = blockIdx.x * blockDim.x + threadIdx.x;
    int pixel_location_j = blockIdx.y * blockDim.y + threadIdx.y;


    if (pixel_location_i >= cuConstRendererParams.imageWidth ||
        pixel_location_j >= cuConstRendererParams.imageHeight)
        return;

    __shared__ float invWidth = 1.f / cuConstRendererParams.imageWidth;
    __shared__ float invHeight = 1.f / cuConstRendererParams.imageHeight;

    float2 pixel_center_Norm = make_float2(
        (pixel_location_i + 0.5f) * invWidth,
        (pixel_location_j + 0.5f) * invHeight);
    float2 pixel_left_down = make_float2(
        pixel_location_i * invWidth,
        pixel_location_j * invHeight);
    float2 pixel_right_up = make_float2(
        (pixel_location_i + 1) * invWidth,
        (pixel_location_j + 1) * invHeight);

    float4 *pixelPtr = (float4 *)(&cuConstRendererParams.imageData[4 * (pixel_location_j * cuConstRendererParams.imageWidth + pixel_location_i)]);

    int index3;
    float3 p;
    float radius;
    int test_intersection;
    float diffX, diffY;
    float pixelDist;
    // Because render object depth is up, so SEQUNCE-DOING traverse
    for (int circle_index = 0; circle_index < cuConstRendererParams.numCircles; circle_index++)
    {
        index3 = 3 * circle_index;
        p = *(float3 *)(&cuConstRendererParams.position[index3]);
        radius = cuConstRendererParams.radius[circle_index];

        // TEST CIRCLE-PIXEL INTERSECTION
        test_intersection = circleInBoxConservative(
            p.x, p.y, radius,
            pixel_left_down.x, pixel_right_up.x,
            pixel_right_up.y, pixel_left_down.y);
        if (test_intersection == 1)
            test_intersection = circleInBox(
                p.x, p.y, radius,
                pixel_left_down.x, pixel_right_up.x,
                pixel_right_up.y, pixel_left_down.y);

        if (test_intersection == 0)
            continue;

        diffX = p.x - pixel_center_Norm.x;
        diffY = p.y - pixel_center_Norm.y;
        pixelDist = diffX * diffX + diffY * diffY;
        shadePixel_Hamster(pixelPtr, circle_index, pixelDist, p);
    }
}

CudaRenderer::CudaRenderer()
{
    image = NULL;

    this->num_Circles = 0;
    this->positions = nullptr;
    this->velocitys = nullptr;
    this->colors = nullptr;
    this->radius = nullptr;

    this->cudaDevicePositions = NULL;
    this->cudaDeviceVelocity = NULL;
    this->cudaDeviceColor = NULL;
    this->cudaDeviceRadius = NULL;
    this->cudaDeviceImageData = NULL;
};

CudaRenderer::~CudaRenderer()
{
    if (image)
        delete image;

    if (positions)
    {
        delete[] positions;
        delete[] velocitys;
        delete[] colors;
        delete[] radius;
    }

    if (cudaDevicePositions)
    {
        cudaFree(cudaDevicePositions);
        cudaFree(cudaDeviceVelocity);
        cudaFree(cudaDeviceColor);
        cudaFree(cudaDeviceRadius);
        cudaFree(cudaDeviceImageData);
    }
}

const Image *CudaRenderer::getImage()
{
    printf("Copy Image from Device to Host\n");

    checkCudaErrors(
        cudaMemcpy(
            image->data,
            cudaDeviceImageData,
            sizeof(float) * image->width * image->height * 4,
            cudaMemcpyDeviceToHost));

    return image;
}

void CudaRenderer::loadScene(SceneName name)
{
    this->sceneName = name;
    loadCircleScene(name, this->num_Circles,
                    this->positions, this->velocitys, this->colors, this->radius);
}

void CudaRenderer::setup()
{
    int deviceCount = 0;
    std::string name;
    auto err = cudaGetDeviceCount(&deviceCount);

    printf("---------------------------------------------------------\n");
    printf("Initializing CUDA for CudaRenderer\n");
    printf("Found %d CUDA devices\n", deviceCount);

    for (int i = 0; i < deviceCount; i++)
    {
        cudaDeviceProp deviceProps;
        checkCudaErrors(cudaGetDeviceProperties(&deviceProps, i));
        name = deviceProps.name;

        printf("Device %d: %s\n", i, deviceProps.name);
        printf("   SMs:        %d\n", deviceProps.multiProcessorCount);
        printf("   Global mem: %.0f MB\n", static_cast<float>(deviceProps.totalGlobalMem) / (1024 * 1024));
        printf("   CUDA Cap:   %d.%d\n", deviceProps.major, deviceProps.minor);
        printf("   MaxThreadsPerBlock : %d\n", deviceProps.maxThreadsPerBlock);
    }

    printf("---------------------------------------------------------\n");

    cudaMalloc(&cudaDevicePositions, sizeof(float) * 3 * num_Circles);
    cudaMalloc(&cudaDeviceVelocity, sizeof(float) * 3 * num_Circles);
    cudaMalloc(&cudaDeviceColor, sizeof(float) * 3 * num_Circles);
    cudaMalloc(&cudaDeviceRadius, sizeof(float) * num_Circles);
    cudaMalloc(&cudaDeviceImageData, sizeof(float) * 4 * image->width * image->height);

    cudaMemcpy(cudaDevicePositions, positions, sizeof(float) * 3 * num_Circles, cudaMemcpyHostToDevice);
    cudaMemcpy(cudaDeviceVelocity, velocitys, sizeof(float) * 3 * num_Circles, cudaMemcpyHostToDevice);
    cudaMemcpy(cudaDeviceColor, colors, sizeof(float) * 3 * num_Circles, cudaMemcpyHostToDevice);
    cudaMemcpy(cudaDeviceRadius, radius, sizeof(float) * num_Circles, cudaMemcpyHostToDevice);

    GlobalConstStruct params;
    params.sceneName = sceneName;
    params.num_Circles = num_Circles;
    params.imageWidth = image->width;
    params.imageHeight = image->height;
    params.position = cudaDevicePositions;
    params.velocity = cudaDeviceVelocity;
    params.color = cudaDeviceColor;
    params.radius = cudaDeviceRadius;
    params.imageData = cudaDeviceImageData;
    cudaMemcpyToSymbol(cuConstRenderParams, &params, sizeof(GlobalConstStruct));

    int *permX;
    int *permY;
    float *value1D;
    getNoiseTables(&permX, &permY, &value1D);
    cudaMemcpyToSymbol(cuConstNoiseXPermutationTable, permX, sizeof(int) * 256);
    cudaMemcpyToSymbol(cuConstNoiseYPermutationTable, permY, sizeof(int) * 256);
    cudaMemcpyToSymbol(cuConstNoise1DValueTable, value1D, sizeof(float) * 256);

    // last, copy over the color table that's used by the shading
    // function for circles in the snowflake demo

    float lookupTable[COLOR_MAP_SIZE][3] = {
        {1.f, 1.f, 1.f},
        {1.f, 1.f, 1.f},
        {.8f, .9f, 1.f},
        {.8f, .9f, 1.f},
        {.8f, 0.8f, 1.f},
    };

    cudaMemcpyToSymbol(cuConstColorRamp, lookupTable, sizeof(float) * 3 * COLOR_MAP_SIZE);
}

void CudaRenderer::allocOutputImage(int width, int height)
{
    if (image)
        delete image;
    image = new Image(width, height);
}

void CudaRenderer::clearImage()
{
    dim3 blockDim(16, 16, 1);
    dim3 gridDim(
        (image->width + blockDim.x - 1) / blockDim.x,
        (image->height + blockDim.y - 1) / blockDim.y);

    if (sceneName == SNOWFLAKES || sceneName == SNOWFLAKES_SINGLE_FRAME)
    {
        kernelClearImageSnowflake<<<gridDim, blockDim>>>();
    }
    else
    {
        kernelClearImage<<<gridDim, blockDim>>>(1.f, 1.f, 1.f, 1.f);
    }
    checkCudaErrors(cudaDeviceSynchronize());
}

void CudaRenderer::advanceAnimation()
{
    // 256 threads per block is a healthy number
    dim3 blockDim(256, 1);
    dim3 gridDim((num_Circles + blockDim.x - 1) / blockDim.x);

    // only the snowflake scene has animation
    if (sceneName == SNOWFLAKES)
    {
        kernelAdvanceSnowflake<<<gridDim, blockDim>>>();
    }
    else if (sceneName == BOUNCING_BALLS)
    {
        kernelAdvanceBouncingBalls<<<gridDim, blockDim>>>();
    }
    else if (sceneName == HYPNOSIS)
    {
        kernelAdvanceHypnosis<<<gridDim, blockDim>>>();
    }
    else if (sceneName == FIREWORKS)
    {
        kernelAdvanceFireWorks<<<gridDim, blockDim>>>();
    }
    checkCudaErrors(cudaDeviceSynchronize());
}

void CudaRenderer::render()
{
   if (this->sceneName == SceneName::SNOWFLAKES || this->sceneName == SceneName::SNOWFLAKES_SINGLE_FRAME)
    {
        dim3 blockDim(256, 1);
        dim3 gridDim((numCircles + blockDim.x - 1) / blockDim.x);
        kernelRenderCircles<<<gridDim, blockDim>>>();
        checkCudaErrors(cudaDeviceSynchronize());
    }
    else
    {
        dim3 blockDim(16, 16, 1);
        dim3 gridDim((image->width + blockDim.x - 1) / blockDim.x,
                     (image->height + blockDim.y - 1) / blockDim.y);
        kernelRenderPixel<<<gridDim, blockDim>>>();
        auto err = cudaGetLastError();
        if (err != cudaSuccess)
        {
            std::cout << "Cuda error: " << cudaGetErrorString(err) << std::endl;
        }
        checkCudaErrors(cudaDeviceSynchronize());
    }
}